#!/usr/bin/env python3
"""
快速训练安全帽检测模型
简化版训练脚本，适合快速开始训练
"""

from ultralytics import YOLO
import torch

def main():
    print("=" * 50)
    print("快速训练安全帽检测模型")
    print("=" * 50)
    
    # 检查设备
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"使用设备: {device}")
    
    if torch.cuda.is_available():
        print(f"GPU: {torch.cuda.get_device_name(0)}")
    
    # 加载预训练模型
    print("\n加载YOLOv8预训练模型...")
    model = YOLO('yolov8n.pt')  # nano版本，训练速度快
    
    # 开始训练
    print("\n开始训练...")
    results = model.train(
        data='data.yaml',           # 数据集配置文件
        epochs=50,                  # 训练轮数 (减少到50轮用于快速验证)
        imgsz=640,                  # 输入图像大小
        batch=8,                    # 批次大小 (较小批次，适合内存有限的情况)
        device=device,              # 自动选择设备
        workers=2,                  # 数据加载线程数
        project='runs/detect',      # 项目保存目录
        name='helmet_detection',    # 实验名称
        save=True,                  # 保存检查点
        save_period=10,             # 每10轮保存一次
        patience=15,                # 早停耐心值
        cache=False,                # 不缓存图像到内存
        augment=True,               # 使用数据增强
        verbose=True                # 详细输出
    )
    
    print("\n训练完成!")
    print(f"最佳模型保存在: runs/detect/helmet_detection/weights/best.pt")
    print(f"最新模型保存在: runs/detect/helmet_detection/weights/last.pt")
    
    # 验证模型
    print("\n验证模型...")
    metrics = model.val()
    print(f"mAP50: {metrics.box.map50:.3f}")
    print(f"mAP50-95: {metrics.box.map:.3f}")

if __name__ == "__main__":
    main() 